import matplotlib.pyplot as plt
import torch
from torch import nn
import torchvision

from torch.utils.data import DataLoader
from torchvision import datasets, transforms

import data_setup
from pathlib import Path
from typing import Dict, List
from tqdm.auto import tqdm

from torch.utils.tensorboard import SummaryWriter

from engine import train_step, test_step
from utils import save_model

# Set seeds
def set_seeds(seed: int=42):
    """Sets random sets for torch operations.
    Args:
        seed (int, optional): Random seed to set. Defaults to 42.
    """
    # Set the seed for general torch operations
    torch.manual_seed(seed)
    # Set the seed for CUDA torch operations (ones that happen on the GPU)
    torch.cuda.manual_seed(seed)


def create_writer(experiment_name: str,
                  model_name: str,
                  extra: str = None) -> torch.utils.tensorboard.writer.SummaryWriter():
    """Creates a torch.utils.tensorboard.writer.SummaryWriter() instance saving to a specific log_dir.

    log_dir is a combination of runs/timestamp/experiment_name/model_name/extra.

    Where timestamp is the current date in YYYY-MM-DD format.

    Args:
        experiment_name (str): Name of experiment.
        model_name (str): Name of model.
        extra (str, optional): Anything extra to add to the directory. Defaults to None.

    Returns:
        torch.utils.tensorboard.writer.SummaryWriter(): Instance of a writer saving to log_dir.

    Example usage:
        # Create a writer saving to "runs/2022-06-04/data_10_percent/effnetb2/5_epochs/"
        writer = create_writer(experiment_name="data_10_percent",
                               model_name="effnetb2",
                               extra="5_epochs")
        # The above is the same as:
        writer = SummaryWriter(log_dir="runs/2022-06-04/data_10_percent/effnetb2/5_epochs/")
    """
    from datetime import datetime
    import os

    # Get timestamp of current date (all experiments on certain day live in same folder)
    timestamp = datetime.now().strftime("%Y-%m-%d")  # returns current date in YYYY-MM-DD format

    if extra:
        # Create log directory path
        log_dir = os.path.join("runs", timestamp, experiment_name, model_name, extra)
    else:
        log_dir = os.path.join("runs", timestamp, experiment_name, model_name)

    print(f"[INFO] Created SummaryWriter, saving to: {log_dir}...")
    return SummaryWriter(log_dir=log_dir)

def train(model: torch.nn.Module,
          train_dataloader: torch.utils.data.DataLoader,
          test_dataloader: torch.utils.data.DataLoader,
          optimizer: torch.optim.Optimizer,
          loss_fn: torch.nn.Module,
          epochs: int,
          device: torch.device,
          writer: torch.utils.tensorboard.writer.SummaryWriter # new
          ) -> Dict[str, List]:

    # Create empty results dictionary
    results = {"train_loss": [],
               "train_acc": [],
               "test_loss": [],
               "test_acc": []
    }

    # Loop through training and testing steps for a number of epochs
    for epoch in tqdm(range(epochs)):
        train_loss, train_acc = train_step(model=model,
                                          dataloader=train_dataloader,
                                          loss_fn=loss_fn,
                                          optimizer=optimizer,
                                          device=device)
        test_loss, test_acc = test_step(model=model,
          dataloader=test_dataloader,
          loss_fn=loss_fn,
          device=device)

        # Print out what's happening
        print(
          f"Epoch: {epoch+1} | "
          f"train_loss: {train_loss:.4f} | "
          f"train_acc: {train_acc:.4f} | "
          f"test_loss: {test_loss:.4f} | "
          f"test_acc: {test_acc:.4f}"
        )

        # Update results dictionary
        results["train_loss"].append(train_loss)
        results["train_acc"].append(train_acc)
        results["test_loss"].append(test_loss)
        results["test_acc"].append(test_acc)

        # See if there's a writer, if so, log to it
        if writer:
            # Add results to SummaryWriter
            writer.add_scalars(main_tag="Loss",
                               tag_scalar_dict={"train_loss": train_loss,
                                                "test_loss": test_loss},
                               global_step=epoch)
            writer.add_scalars(main_tag="Accuracy",
                               tag_scalar_dict={"train_acc": train_acc,
                                                "test_acc": test_acc},
                               global_step=epoch)

            # Close the writer
            writer.close()
        else:
            pass

    return results


# Create an EffNetB0 feature extractor
def create_effnetb0():
    # 1. Get the base mdoel with pretrained weights and send to target device
    weights = torchvision.models.EfficientNet_B0_Weights.DEFAULT
    model = torchvision.models.efficientnet_b0(weights=weights).to(device)

    # 2. Freeze the base model layers
    for param in model.features.parameters():
        param.requires_grad = False

    # 3. Set the seeds
    set_seeds()

    # 4. Change the classifier head
    model.classifier = nn.Sequential(
        nn.Dropout(p=0.2),
        nn.Linear(in_features=1280, out_features=OUT_FEATURES)
    ).to(device)

    # 5. Give the model a name
    model.name = "effnetb0"
    print(f"[INFO] Created new {model.name} model.")
    return model

# Create an EffNetB2 feature extractor
def create_effnetb2():
    # 1. Get the base model with pretrained weights and send to target device
    weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
    model = torchvision.models.efficientnet_b2(weights=weights).to(device)

    # 2. Freeze the base model layers
    for param in model.features.parameters():
        param.requires_grad = False

    # 3. Set the seeds
    set_seeds()

    # 4. Change the classifier head
    model.classifier = nn.Sequential(
        nn.Dropout(p=0.3),
        nn.Linear(in_features=1408, out_features=OUT_FEATURES)
    ).to(device)

    # 5. Give the model a name
    model.name = "effnetb2"
    print(f"[INFO] Created new {model.name} model.")
    return model


def try_gpu(i=0):
    """如果存在，则返回gpu(i)，否则返回cpu()"""
    if torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')


def load_data(BATCH_SIZE = 32):
    data_path = Path("dataset/")
    # Setup directories
    data_10_percent_path = data_path / "pizza_steak_sushi"
    data_20_percent_path = data_path / "pizza_steak_sushi_20_percent"
    # Setup training directory paths
    train_dir_10_percent = data_10_percent_path / "train"
    train_dir_20_percent = data_20_percent_path / "train"
    # Setup testing directory paths (note: use the same test dataset for both to compare the results)
    test_dir = data_10_percent_path / "test"

    # Setup ImageNet normalization levels (turns all images into similar distribution as ImageNet)
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])

    # Compose transforms into a pipeline
    simple_transform = transforms.Compose([
        transforms.Resize((224, 224)), # 1. Resize the images
        transforms.ToTensor(), # 2. Turn the images into tensors with values between 0 & 1
        normalize # 3. Normalize the images so their distributions match the ImageNet dataset
    ])

    # Create 10% training and test DataLoaders
    train_dataloader_10_percent, test_dataloader, class_names = data_setup.create_dataloaders(train_dir=train_dir_10_percent,
        test_dir=test_dir,
        transform=simple_transform,
        batch_size=BATCH_SIZE
    )

    # Create 20% training and test data DataLoders
    train_dataloader_20_percent, test_dataloader, class_names = data_setup.create_dataloaders(train_dir=train_dir_20_percent,
        test_dir=test_dir,
        transform=simple_transform,
        batch_size=BATCH_SIZE
    )

    # Find the number of samples/batches per dataloader (using the same test_dataloader for both experiments)
    print(f"Number of batches of size {BATCH_SIZE} in 10 percent training data: {len(train_dataloader_10_percent)}")
    print(f"Number of batches of size {BATCH_SIZE} in 20 percent training data: {len(train_dataloader_20_percent)}")
    print(f"Number of batches of size {BATCH_SIZE} in testing data: {len(train_dataloader_10_percent)} (all experiments will use the same test set)")
    print(f"Number of classes: {len(class_names)}, class names: {class_names}")

    return train_dataloader_10_percent, train_dataloader_20_percent, test_dataloader, class_names





if __name__ == '__main__':
    train_dataloader_10_percent, train_dataloader_20_percent, test_dataloader, class_names = load_data()
    # Get num out features (one for each class pizza, steak, sushi)
    OUT_FEATURES = len(class_names)
    device = try_gpu()
    print(device)

    effnetb0 = create_effnetb0()
    effnetb2 = create_effnetb2()

    # 1. Create epochs list
    num_epochs = [5, 10]
    # 2. Create models list (need to create a new model for each experiment)
    models = ["effnetb0", "effnetb2"]
    # 3. Create dataloaders dictionary for various dataloaders
    train_dataloaders = {"data_10_percent": train_dataloader_10_percent,
                         "data_20_percent": train_dataloader_20_percent}


    # 1. Set the random seeds
    set_seeds()

    # 2. Keep track of experiment numbers
    experiment_number = 0

    # 3. Loop through each DataLoader
    for dataloader_name, train_dataloader in train_dataloaders.items():

        # 4. Loop through each number of epochs
        for epochs in num_epochs:

            # 5. Loop through each model name and create a new model based on the name
            for model_name in models:

                # 6. Create information print outs
                experiment_number += 1
                print(f"[INFO] Experiment number: {experiment_number}")
                print(f"[INFO] Model: {model_name}")
                print(f"[INFO] DataLoader: {dataloader_name}")
                print(f"[INFO] Number of epochs: {epochs}")

                # 7. Select the model
                if model_name == "effnetb0":
                    model = create_effnetb0()  # creates a new model each time (important because we want each experiment to start from scratch)
                else:
                    model = create_effnetb2()  # creates a new model each time (important because we want each experiment to start from scratch)

                # 8. Create a new loss and optimizer for every model
                loss_fn = nn.CrossEntropyLoss()
                optimizer = torch.optim.Adam(params=model.parameters(), lr=0.001)

                # 9. Train target model with target dataloaders and track experiments
                train(model=model,
                      train_dataloader=train_dataloader,
                      test_dataloader=test_dataloader,
                      optimizer=optimizer,
                      loss_fn=loss_fn,
                      epochs=epochs,
                      device=device,
                      writer=create_writer(experiment_name=dataloader_name,
                                           model_name=model_name,
                                           extra=f"{epochs}_epochs"))

                # 10. Save the model to file so we can get back the best model
                save_filepath = f"07_{model_name}_{dataloader_name}_{epochs}_epochs.pth"
                save_model(model=model,
                           target_dir="models",
                           model_name=save_filepath)
                print("-" * 50 + "\n")